import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse
import os
import seaborn as sns

def parse_args():
    parser = argparse.ArgumentParser(description="Generate Faceted Bar Plot.")
    parser.add_argument('--files','-f', nargs='+', required=True, help='List of files (CSV or Excel) to analyze')
    parser.add_argument('--exclude', nargs='+', default=['BART'], help='Methods to exclude (default: BART)')
    parser.add_argument('--output','-o', default='faceted_bar_plot.png', help='Output filename for the plot')
    return parser.parse_args()

def read_data(file_path):
    _, ext = os.path.splitext(file_path)
    try:
        if ext.lower() in ['.xlsx', '.xls']:
            return pd.read_excel(file_path)
        else:
            return pd.read_csv(file_path)
    except Exception as e:
        print(f"Error reading {file_path}: {e}")
        return None

def get_aggregated_stats(files, exclude_methods):
    """
    读取文件，合并数据。
    返回的数据结构：包含 Sample (文件名), Method, Metric, Mean, Std
    """
    all_data = []
    
    for f in files:
        df = read_data(f)
        if df is not None and not df.empty:
            # 1. 过滤方法
            df = df[~df['method'].isin(exclude_methods)].copy()
            
            # 2. 重命名
            df['method'] = df['method'].replace('ours(MTM)', 'ours(deepsignal3)')
            
            # 3. 过滤 Coverage (保留 10 的整数倍)
            # 这一步是为了计算标准差时排除杂项数据的影响
            if 'coverage' in df.columns:
                df['coverage'] = pd.to_numeric(df['coverage'], errors='coerce')
                df = df.dropna(subset=['coverage'])
                df = df[np.isclose(df['coverage'] % 10, 0)]
            
            # 提取样本名称 (去掉后缀)
            sample_name = os.path.basename(f).split('.')[0]
            # 如果文件名像 "HG001.xlsx - Sheet1.csv"，再处理一下
            if " - " in sample_name:
                sample_name = sample_name.split(" - ")[0]
            if sample_name.endswith("xlsx"):
                 sample_name = sample_name.replace(".xlsx", "")
            
            df['Sample'] = sample_name
            
            cols = ['Sample', 'method', 'pearson', 'rsquare', 'spearman', 'RMSE']
            if all(c in df.columns for c in cols[2:]):
                all_data.append(df[cols])
            
    if not all_data:
        return None
        
    combined_df = pd.concat(all_data, ignore_index=True)
    return combined_df

def plot_faceted_bar(df, output_file):
    # 将宽格式数据转换为长格式，方便 Seaborn 分面绘图
    # id_vars: Sample, method
    # value_vars: pearson, rsquare, spearman, RMSE
    melted_df = df.melt(id_vars=['Sample', 'method'], 
                        value_vars=['pearson', 'rsquare', 'spearman', 'RMSE'],
                        var_name='Metric', value_name='Score')

    # 设置风格
    sns.set_style("whitegrid")
    sns.set_context("paper", font_scale=1.4)
    
    # 定义颜色映射
    palette = {
        'Dorado': '#1f77b4',       # Blue
        'ours(deepsignal3)': '#d62728', # Red (Highlight)
        'Rockfish': '#2ca02c',     # Green
        'DeepMod2': '#ff7f0e'      # Orange
    }
    
    # 按照 Sample 排序 (HG001, HG002...)
    samples = sorted(melted_df['Sample'].unique())
    
    # 创建分面绘图 (Catplot)
    # col="Metric": 按指标分列/分面
    # col_wrap=2: 每行2个图 -> 2x2 布局
    # sharey=False: 每个指标的Y轴范围独立 (因为RMSE是0.1，Pearson是0.9)
    g = sns.catplot(
        data=melted_df, 
        kind="bar",
        x="Sample", 
        y="Score", 
        hue="method",
        col="Metric", 
        col_wrap=2,
        palette=palette,
        height=5, 
        aspect=1.5,
        sharey=False,
        errorbar='sd', # 显示标准差作为误差线
        capsize=0.1,   # 误差线帽子大小
        err_kws={'linewidth': 1.5, 'color': '0.2'}, # 误差线样式
        legend_out=True # 图例放在外面
    )
    
    # 自定义每个子图的标题和标签
    titles = {
        'pearson': 'Pearson Correlation (Higher is Better)',
        'rsquare': 'R-Square (Higher is Better)',
        'spearman': 'Spearman Correlation (Higher is Better)',
        'RMSE': 'RMSE (Lower is Better)'
    }
    
    for ax, col_name in zip(g.axes.flat, g.col_names):
        # 设置标题
        ax.set_title(titles.get(col_name, col_name), fontweight='bold')
        
        # 优化 X 轴标签
        ax.set_xlabel('')
        
        # 针对 RMSE 可能需要的特殊处理 (如反转Y轴? 一般柱状图不反转，低的柱子短就是好)
        if col_name == 'RMSE':
            # 可以在图里标注 "Lower is better"
            pass
            
    # 调整图例
    # Seaborn catplot 的图例有时很难控制，我们重新设置
    sns.move_legend(g, "upper center", bbox_to_anchor=(0.5, -0.05), ncol=4, title=None, frameon=False)
    
    plt.tight_layout()
    # 留出底部空间给图例
    plt.subplots_adjust(bottom=0.12)
    
    print(f"Saving faceted bar plot to {output_file}")
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.show()

def main():
    args = parse_args()
    
    print("Reading and aggregating data...")
    df = get_aggregated_stats(args.files, args.exclude)
    
    if df is not None:
        plot_faceted_bar(df, args.output)
    else:
        print("No valid data found.")

if __name__ == "__main__":
    main()